# 将 Markdown 文档内容向量化并存储到 Milvus 向量数据库
# Markdown 文档处理 1、扫描指定目录（D:\PyCharm...\中间结果2）下的所有 .md 文件
#   2、使用 MarkdownHeaderTextSplitter 按标题层级（#, ##, ###）分割文档内容，保留标题作为元数据。
# 向量化存储
#   1、使用 Ollama 嵌入模型（nomic-embed-text）将文本转换为 768 维向量。
#   2、将向量和文本数据存入 Milvus 数据库，包含以下字段：
#       text：分割后的文本内容
#       header_1/2/3：各级标题（元数据）
#       metadata：原始元数据（JSON）
#       vector：向量数据
# 语义检索 1、提供交互式查询接口，输入自然语言问题后，返回最相关的文档片段（基于向量相似度）。

from config import envConfig, config
import os
from langchain_milvus import Milvus
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_ollama import OllamaEmbeddings
from langchain.retrievers import EnsembleRetriever
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import logging
from pymilvus import utility

logger = logging.getLogger('vectorize')

MAX_CHUNK_SIZE = 400
CHUNK_OVERLAP = 100
embedding = OllamaEmbeddings(
    model="nomic-embed-text",
    base_url=envConfig.OLLAMA_URL
)


def get_all_md():
    # 目标目录路径所有.md文件
    target_dir = r"D:\PyCharmMiscProject\pythonProjectYA\中间结果\中间结果2"

    # 路径合法性
    if not os.path.exists(target_dir):
        raise FileNotFoundError(f"目录不存在: {target_dir}")
    if not os.path.isdir(target_dir):
        raise NotADirectoryError(f"路径不是目录: {target_dir}")

    # 获取目录下所有 .md 文件的完整路径
    md_files = []
    for filename in os.listdir(target_dir):
        if filename.endswith(".md"):
            full_path = os.path.join(target_dir, filename)
            md_files.append(full_path)

    return md_files


from pathlib import Path

def get_all_md_perfect():
    # 用 pathlib 库更优雅地处理路径
    target_dir = Path(r"D:\PyCharmMiscProject\pythonProjectYA\中间结果\中间结果2")
    if not target_dir.is_dir():
        raise NotADirectoryError(f"路径不是目录: {target_dir}")
    return [str(p) for p in target_dir.glob("*.md")]


def split_by_md_header(mdDocs):
    """
    Split with metadata like "Metadata (("Header_1", "xxx"), ("Header_2", "xxxx")) page_content="xxx..."
    :param mdDocs:
    :return:
    """
    headers_to_split_on = [("#", "Header_1"), ("##", "Header_2"), ("###", "Header_3")]
    md_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on, return_each_line=True)
    md_header_splits = md_splitter.split_text(mdDocs)
    final_splits = []
    for _, i in enumerate(md_header_splits):
        print(f"Index: {_}, Metadata: {i.metadata} Content: {i.page_content}")
        print("---------------------")
    return md_header_splits


def get_vector_store(collectionName: str = 'LangChainCollection', is_local=True):
    """
    获取 Milvus 向量存储连接（修复版）
    """
    # 连接参数配置
    if is_local:
        milvusConn = {
            "uri": envConfig.MILVUS_LOCAL_URI,
            "user": envConfig.MILVUS_USER,
            "password": envConfig.MILVUS_PASSWORD
        }
    else:
        milvusConn = {
            "uri": envConfig.BACKUP_MILVUS_URI,
            "user": envConfig.BACKUP_MILVUS_USERNAME,
            "password": envConfig.BACKUP_MILVUS_PASSWORD,
            "secure": True
        }

    # 建立连接
    try:
        connections.connect(**milvusConn)
        logger.info(f"成功连接 Milvus: {milvusConn['uri']}")
    except Exception as e:
        logger.error(f"Milvus 连接失败: {str(e)}")
        raise

    # 定义集合 Schema
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=False),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=4096),
        FieldSchema(name="header_1", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY'),
        FieldSchema(name="header_2", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY'),
        FieldSchema(name="header_3", dtype=DataType.VARCHAR, max_length=256, default_value='EMPTY'),
        FieldSchema(name="metadata", dtype=DataType.JSON),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=768)
    ]

    schema = CollectionSchema(fields, description="LangChain Collection")

    # 检查集合是否存在
    if utility.has_collection(collectionName):
        logger.warning(f"集合 {collectionName} 已存在，正在删除...")
        utility.drop_collection(collectionName)

    # 创建新集合
    collection = Collection(name=collectionName, schema=schema)

    # 创建索引
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}
    }
    collection.create_index(field_name="vector", index_params=index_params)
    collection.load()

    return collection


def vectorize_all(drop_old=True):
    print("Starting to vectorize all markdown files")

    # 初始化 Milvus 连接（通过 get_vector_store 自动处理清理）
    collection_name = config.APP_NAME.replace(' ', '_')
    vector_store = get_vector_store(
        collectionName=collection_name,
        is_local=True
    )

    # 生成文档数据
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url=envConfig.OLLAMA_URL
    )
    files = get_all_md_perfect()
    all_splits = []
    for file in files:
        with open(file, 'r', encoding='utf-8') as f:
            mdDocs = f.read()
        all_splits.extend(split_by_md_header(mdDocs))

    # 准备插入数据
    documents = [{
        "id": idx,
        "text": split.page_content,
        "header_1": split.metadata.get("Header_1", ""),
        "header_2": split.metadata.get("Header_2", ""),
        "header_3": split.metadata.get("Header_3", ""),
        "metadata": split.metadata,
        "vector": embeddings.embed_query(split.page_content)
    } for idx, split in enumerate(all_splits)]

    # 插入数据
    vector_store.insert(documents)
    print(f"数据已存入集合: {collection_name}")



# vector = get_vector_store(embedding, "Intention_Instruction")
# vector.as_retriever().invoke("指令类型")
def test_vectorize():
    """
    测试向量存储功能（优先使用 Docker 实例）
    """
    # 直接使用 Docker 实例配置
    milvusConn = {
        "uri": envConfig.MILVUS_LOCAL_URI,
        "user": envConfig.MILVUS_USER,
        "password": envConfig.MILVUS_PASSWORD
    }

    # 移除错误的 drop_old 逻辑（集合管理应该在 vectorize_all 中完成）
    embeddings = OllamaEmbeddings(
        model="nomic-embed-text",
        base_url=envConfig.OLLAMA_URL
    )

    retriever = Milvus(
        collection_name=config.APP_NAME.replace(' ', '_'),
        connection_args=milvusConn,
        embedding_function=embeddings
    ).as_retriever(search_kwargs={"k": 1})  # 返回结果的数量

    is_over = False
    while not is_over:
        query = input("请输入查询，回车查询exit退出：")
        if query == "exit":
            is_over = True
            continue
        result = retriever.get_relevant_documents(query)  # 直接使用单一 retriever
        for i, ele in enumerate(result):
            print(f"{i} - Found Doc: {ele.metadata}, {ele.page_content}")

vectorize_all()
test_vectorize()
print("FINISH!")